#include <errno.h>
#include <bpf/bpf.h>
#include <bpf/libbpf.h>
#include <stdio.h>
#include <assert.h>
#include <arpa/inet.h>
#include <netinet/in.h>
#include <sys/socket.h>
#include <linux/if_ether.h>
#include <linux/ip.h>
#include <linux/ipv6.h>
#include <linux/tcp.h>
#include "xdp_test.skel.h"
#include <asm/unistd.h>


#define PACKET_SIZE 100

struct xdp_bpf *skel;
int prog_fd;

static inline int sys_bpf(enum bpf_cmd cmd, union bpf_attr *attr,
			  unsigned int size)
{
	return syscall(__NR_bpf, cmd, attr, size);
}

int bpf_prog_test_run_opts_new(int prog_fd, struct bpf_test_run_opts *opts)
{
	union bpf_attr attr;
	int ret;

	memset(&attr, 0, sizeof(attr));
	attr.test.prog_fd = prog_fd;
	attr.test.cpu = opts->cpu;
	attr.test.flags = opts->flags;
	attr.test.repeat = opts->repeat;
	attr.test.duration = opts->duration;
	attr.test.ctx_size_in = opts->ctx_size_in;
	attr.test.ctx_size_out = opts->ctx_size_out;
	attr.test.data_size_in = opts->data_size_in;
	attr.test.data_size_out = opts->data_size_out;
	attr.test.ctx_in = opts->ctx_in;
	attr.test.ctx_out = opts->ctx_out;
	attr.test.data_in = opts->data_in;
	attr.test.data_out = opts->data_out;

	ret = sys_bpf(BPF_PROG_TEST_RUN, &attr, sizeof(attr));

    opts->data_size_out = attr.test.data_size_out;
	opts->ctx_size_out = attr.test.ctx_size_out;
	opts->duration = attr.test.duration;
	opts->retval = attr.test.retval;

	return ret;
}



void test_init() {
    skel = xdp_bpf__open_and_load();
    assert(skel != NULL);
    prog_fd = bpf_program__fd(skel->progs.xdp_shutdown);
}

void test_finish() {
    xdp_bpf__destroy(skel);
}

static int run_xdp_test(void *packet, size_t size) {
    struct xdp_md ctx = {
        .data = (uint32_t)packet,
        .data_end = (uint32_t)(packet + size),
    };

    struct bpf_test_run_opts opts = {
        .sz = sizeof(struct bpf_test_run_opts),
        .ctx_in = &ctx,
        .ctx_size_in = sizeof(ctx),
        .ctx_out = &ctx,
        .ctx_size_out = sizeof(ctx)
    };

    return bpf_prog_test_run_opts_new(prog_fd, &opts);
}

void test_packet_parsing() {
    unsigned char packet[PACKET_SIZE] = {0};
    struct ethhdr *eth = (struct ethhdr *)packet;
    struct iphdr *ip = (struct iphdr *)(packet + sizeof(struct ethhdr));
    struct tcphdr *tcp = (struct tcphdr *)(packet + sizeof(struct ethhdr) + sizeof(struct iphdr));

    eth->h_proto = htons(ETH_P_IP);
    ip->version = 4;
    ip->ihl = 5;
    ip->protocol = IPPROTO_TCP;
    tcp->source = htons(12345);
    tcp->dest = htons(80);

    int err = run_xdp_test(packet, PACKET_SIZE);
    printf("err = %d\n", err);
    assert(err == 0);
}

void test_ip_version_check() {
    unsigned char packet[PACKET_SIZE] = {0};
    struct ethhdr *eth = (struct ethhdr *)packet;
    struct iphdr *ip = (struct iphdr *)(packet + sizeof(struct ethhdr));

    eth->h_proto = htons(ETH_P_IP);
    ip->version = 5;  // Invalid IP version

    int err = run_xdp_test(packet, PACKET_SIZE);
    assert(err == 0);
}

void test_tuple_extraction() {
    unsigned char packet[PACKET_SIZE] = {0};
    struct ethhdr *eth = (struct ethhdr *)packet;
    struct iphdr *ip = (struct iphdr *)(packet + sizeof(struct ethhdr));
    struct tcphdr *tcp = (struct tcphdr *)(packet + sizeof(struct ethhdr) + sizeof(struct iphdr));

    eth->h_proto = htons(ETH_P_IP);
    ip->version = 4;
    ip->ihl = 5;
    ip->protocol = IPPROTO_TCP;
    ip->saddr = inet_addr("192.168.1.1");
    ip->daddr = inet_addr("192.168.1.2");
    tcp->source = htons(12345);
    tcp->dest = htons(80);

    int err = run_xdp_test(packet, PACKET_SIZE);
    assert(err == 0);
}

void test_connection_shutdown() {
    unsigned char packet[PACKET_SIZE] = {0};
    struct ethhdr *eth = (struct ethhdr *)packet;
    struct iphdr *ip = (struct iphdr *)(packet + sizeof(struct ethhdr));
    struct tcphdr *tcp = (struct tcphdr *)(packet + sizeof(struct ethhdr) + sizeof(struct iphdr));

    eth->h_proto = htons(ETH_P_IP);
    ip->version = 4;
    ip->ihl = 5;
    ip->protocol = IPPROTO_TCP;
    ip->saddr = inet_addr("192.168.1.1");
    ip->daddr = inet_addr("192.168.1.2");
    tcp->source = htons(12345);
    tcp->dest = htons(80);
    tcp->syn = 1;

    // Add the connection to the auth map to simulate a connection that should be shut down
    struct bpf_sock_tuple tuple = {
        .ipv4 = {
            .saddr = ip->saddr,
            .daddr = ip->daddr,
            .sport = tcp->source,
            .dport = tcp->dest,
        },
    };
    __u32 value = 1;
    bpf_map_update_elem(bpf_map__fd(skel->maps.map_of_auth), &tuple, &value, BPF_ANY);

    int err = run_xdp_test(packet, PACKET_SIZE);
    assert(err == 0);

    // Check if the connection was shut down (RST flag set)
    assert(tcp->rst == 1);
}

int main() {
    test_init();

    test_packet_parsing();
    test_ip_version_check();
    test_tuple_extraction();
    test_connection_shutdown();

    test_finish();
    printf("All tests passed!\n");
    return 0;
}